import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
import seaborn as sns

sns.set_theme()
# 设置中文字体（Windows 下）
ch_font = FontProperties(fname="C:/Windows/Fonts/simhei.ttf")

# 读取数据
df = pd.read_csv("加权后的分组得分.csv")
df = df[pd.to_numeric(df["year"], errors="coerce").notna()]
df["year"] = df["year"].astype(int)

# 压缩段年份（你给出的分布）
sparse_years = [1920, 1935, 1941, 1944]

# 构建映射表
# 稀疏段从位置 0 开始，每年步长为 1
mapped_years = {}
pos = 0
for year in sorted(set(sparse_years)):
    mapped_years[year] = pos
    pos += 1  # 步长 = 1（压缩）

# 稠密段从 pos 继续，每年步长为 3（拉伸）
for year in range(1955, df["year"].max() + 1):
    mapped_years[year] = pos
    pos += 3  # 步长 = 3（拉伸）

# 映射应用到数据中
df = df[df["year"].isin(mapped_years)]
df["mapped_year"] = df["year"].map(mapped_years)

# 分组配置
type_groups = [
    ("Stimulation_and_Adventure", "Documentary_and_Public"),
    ("joy_and_escapism", "Realism_and_Depth"),
    ("Anxiety_Conflict", "Documentary_and_Public"),
]

# 绘图
fig, axes = plt.subplots(nrows=3, ncols=1, figsize=(14, 12), sharex=True)

for i, (col1, col2) in enumerate(type_groups):
    ax = axes[i]

    # 汇总每年的得分
    yearly1 = df.groupby(["year", "mapped_year"])[col1].sum().reset_index()
    yearly2 = df.groupby(["year", "mapped_year"])[col2].sum().reset_index()

    ax.plot(yearly1["mapped_year"], yearly1[col1], label=col1, linewidth=2)
    ax.plot(yearly2["mapped_year"], yearly2[col2], label=col2, linewidth=2)

    ax.set_ylabel("得分", fontproperties=ch_font)
    ax.set_title(f"{col1} vs {col2}", fontproperties=ch_font)
    ax.grid(True)
    ax.legend(loc="upper left", prop=ch_font)

# 设置x轴为原始年份
# 构建 xticks 和 labels
xtick_positions = []
xtick_labels = []

# 控制标注密度
for year in sorted(mapped_years):
    mapped = mapped_years[year]
    # 压缩段全部标
    if year <= 1945:
        xtick_positions.append(mapped)
        xtick_labels.append(str(year))
    else:
        # 稠密段隔5年标一次
        if year % 5 == 0:
            xtick_positions.append(mapped)
            xtick_labels.append(str(year))

axes[-1].set_xticks(xtick_positions)
axes[-1].set_xticklabels(xtick_labels, rotation=45, fontproperties=ch_font)
axes[-1].set_xlabel("年份（1920-1946 已压缩）", fontproperties=ch_font)

# 总标题
fig.suptitle(
    "电影类型随年份得分趋势（压缩早期年份）", fontsize=16, fontproperties=ch_font
)
plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.show()
